package org.xx.armory.commons;

import org.xx.armory.config.ConfigurationException;

import javax.crypto.KeyGenerator;
import javax.crypto.SecretKey;
import javax.crypto.SecretKeyFactory;
import javax.crypto.spec.DESKeySpec;
import java.io.File;
import java.io.IOException;
import java.io.InputStream;
import java.security.GeneralSecurityException;
import java.security.InvalidKeyException;
import java.security.KeyFactory;
import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.KeyStore;
import java.security.KeyStoreException;
import java.security.NoSuchAlgorithmException;
import java.security.NoSuchProviderException;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.SecureRandom;
import java.security.cert.Certificate;
import java.security.cert.CertificateEncodingException;
import java.security.cert.CertificateFactory;
import java.security.cert.X509Certificate;
import java.security.spec.InvalidKeySpecException;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;
import java.util.Base64;
import java.util.Enumeration;
import java.util.Scanner;

import static org.apache.commons.io.FileUtils.openInputStream;
import static org.apache.commons.io.FilenameUtils.getExtension;
import static org.apache.commons.lang3.StringUtils.trimToEmpty;
import static org.xx.armory.commons.Validators.notEmpty;
import static org.xx.armory.commons.Validators.notNull;

/**
 * 密钥相关的工具类。
 */
public final class KeyUtils {
    private static final SecretKeyFactory DES_KEY_FACTORY = createDESKeyFactory();

    private static final KeyFactory RSA_KEY_FACTORY = createRSAKeyFactory();

    private KeyUtils() {
        throw new UnsupportedOperationException();
    }

    private static SecretKeyFactory createDESKeyFactory() {
        try {
            return SecretKeyFactory.getInstance("DES");
        } catch (NoSuchAlgorithmException ex) {
            throw new IllegalStateException("cannot get DES algorithm", ex);
        }
    }

    private static KeyFactory createRSAKeyFactory() {
        try {
            return KeyFactory.getInstance("RSA");
        } catch (NoSuchAlgorithmException ex) {
            throw new ConfigurationException("cannot get DES algorithm", ex);
        }
    }

    /**
     * 生成{@literal RSA}算法的密钥对。
     *
     * @return 生成的密钥对。
     * @throws NoSuchAlgorithmException
     *         如果不存在指定的加密算法。
     * @throws NoSuchProviderException
     *         如果不存在指定算法的默认提供程序。
     */
    public static KeyPair generateRSAKeys()
            throws NoSuchAlgorithmException, NoSuchProviderException {
        final KeyPairGenerator generator = KeyPairGenerator.getInstance("RSA");

        final SecureRandom random = SecureRandom.getInstance("SHA1PRNG");

        generator.initialize(2048, random);

        return generator.generateKeyPair();
    }

    /**
     * 生成{@literal DES}算法的密钥。
     *
     * @return 生成的密钥对。
     * @throws NoSuchAlgorithmException
     *         如果不存在指定的加密算法。
     * @throws NoSuchProviderException
     *         如果不存在指定算法的默认提供程序。
     */
    public static SecretKey generateDESKey()
            throws NoSuchAlgorithmException, NoSuchProviderException {
        final KeyGenerator generator = KeyGenerator.getInstance("DES");

        final SecureRandom random = SecureRandom.getInstance("SHA1PRNG");

        generator.init(56, random);

        return generator.generateKey();
    }

    /**
     * 根据密钥的已编码内容创建{@literal DES}密钥。
     *
     * @param value
     *         密钥的已编码内容。
     * @return 已创建的密钥。
     * @throws IllegalArgumentException
     *         如果参数{@code value}是{@code null}或者是空数组。
     * @throws InvalidKeyException
     *         如果指定的密钥编码内容不正确。
     * @throws InvalidKeySpecException
     *         如果当前运行环境不支持{@literal DES}格式的密钥。
     */
    public static SecretKey buildDESKey(
            byte[] value
    )
            throws InvalidKeyException, InvalidKeySpecException {
        notEmpty(value, "value");

        return DES_KEY_FACTORY.generateSecret(new DESKeySpec(value));
    }

    /**
     * 根据密钥的已编码内容创建{@literal RSA}的公钥。
     *
     * @param value
     *         密钥的已编码内容。
     * @return 已创建的密钥。
     * @throws IllegalArgumentException
     *         如果参数{@code value}是{@code null}或者是空数组。
     * @throws InvalidKeyException
     *         如果指定的密钥编码内容不正确。
     * @throws InvalidKeySpecException
     *         如果当前运行环境不支持{@literal X.509}格式的密钥。
     */
    public static PublicKey buildRSAPublicKey(
            byte[] value
    )
            throws InvalidKeyException, InvalidKeySpecException {
        notEmpty(value, "value");

        return RSA_KEY_FACTORY.generatePublic(new X509EncodedKeySpec(value));
    }

    /**
     * 根据密钥的已编码内容创建{@literal RSA}的私钥。
     *
     * @param value
     *         密钥的已编码内容。
     * @return 已创建的密钥。
     * @throws InvalidKeyException
     *         如果指定的密钥编码内容不正确。
     * @throws InvalidKeySpecException
     *         如果当前运行环境不支持{@literal DES}密钥。
     */
    public static PrivateKey buildRSAPrivateKey(
            byte[] value
    )
            throws InvalidKeyException, InvalidKeySpecException {
        notEmpty(value, "value");

        return RSA_KEY_FACTORY.generatePrivate(new PKCS8EncodedKeySpec(value));
    }

    /**
     * 从指定的证书存储中找到最后一个可用的别名。
     *
     * @param keyStore
     *         证书存储。
     * @return 最后一个可用的别名。
     * @throws KeyStoreException
     *         如果查找别名时出错。
     * @throws IllegalArgumentException
     *         如果参数{@code keyStore}是{@code null}。
     */
    private static String findLastKeyAlias(
            KeyStore keyStore
    )
            throws KeyStoreException {
        notNull(keyStore, "keyStore");

        for (final Enumeration<String> enu = keyStore.aliases(); enu.hasMoreElements(); ) {
            final String alias = enu.nextElement();
            if (keyStore.isKeyEntry(alias)) {
                return alias;
            }
        }

        return null;
    }

    public static Certificate exportCertificate(
            InputStream stream,
            char[] password
    ) {
        notNull(stream, "stream");
        notEmpty(password, "password");

        try {
            // 加载证书对象
            final KeyStore keyStore = KeyStore.getInstance("PKCS12");
            keyStore.load(stream, password);

            final String alias = findLastKeyAlias(keyStore);
            if (alias == null) {
                throw new IllegalArgumentException("cannot find alias from stream");
            }

            return keyStore.getCertificate(alias);
        } catch (GeneralSecurityException | IOException ex) {
            throw new RuntimeException("cannot dump certificate from PKCS12 cert", ex);
        }
    }

    /**
     * 从{@literal PCKS12}格式的证书文件中导出公钥和私钥
     *
     * @param stream
     *         包含PCKS12格式的证书文件流
     * @param password
     *         私钥保护密码
     * @return 导出的密钥对，包含公钥和私钥。
     * @throws IllegalArgumentException
     *         如果参数{@code stream}或者{@code password}是{@code null}，或者参数{@code password}是空数组。。
     */
    public static KeyPair exportPCKS12Keys(
            InputStream stream,
            char[] password
    ) {
        notNull(stream, "stream");
        notEmpty(password, "password");

        try {
            // 加载证书对象
            final KeyStore keyStore = KeyStore.getInstance("PKCS12");
            keyStore.load(stream, password);

            final String alias = findLastKeyAlias(keyStore);
            if (alias == null) {
                throw new IllegalArgumentException("cannot find alias from stream");
            }

            final PrivateKey privateKey = (PrivateKey) keyStore.getKey(alias, password);
            final PublicKey publicKey = keyStore.getCertificate(alias).getPublicKey();

            return new KeyPair(publicKey, privateKey);
        } catch (GeneralSecurityException | IOException ex) {
            throw new RuntimeException("cannot dump keys from PKCS12 cert", ex);
        }
    }

    /**
     * 从{@literal X.509}格式的证书文件中导出公钥
     *
     * @param stream
     *         包含X.509格式的证书文件流
     * @return 导出的公钥。
     * @throws IllegalArgumentException
     *         如果参数{@code stream}是{@code null}。
     */
    public static PublicKey exportX509Keys(
            InputStream stream
    ) {
        notNull(stream, "stream");

        try {
            final CertificateFactory cf = CertificateFactory.getInstance("X.509"); // 指定证书类型
            final X509Certificate cert = (X509Certificate) cf.generateCertificate(stream); // 获取证书

            return cert.getPublicKey();
        } catch (GeneralSecurityException ex) {
            throw new RuntimeException("cannot dump keys from X.509 cert", ex);
        }
    }

    public static String dumpCertificate(
            Certificate certificate
    )
            throws CertificateEncodingException {
        notNull(certificate, "certificate");

        final byte[] encodedKey = certificate.getEncoded();
        return "---- Certificate ----\n" +
                "type: " + certificate.getType() + "\n" +
                "size: " + encodedKey.length + "\n" +
                "encoded: " + Base64.getEncoder().encodeToString(encodedKey);
    }

    public static String dumpPublicKey(
            PublicKey publicKey
    ) {
        notNull(publicKey, "publicKey");

        final byte[] encodedKey = publicKey.getEncoded();
        return "---- Public Key ----\n" +
                "algorithm: " + publicKey.getAlgorithm() + "\n" +
                "format: " + publicKey.getFormat() + "\n" +
                "size: " + encodedKey.length + "\n" +
                "encoded: " + Base64.getEncoder().encodeToString(encodedKey);
    }

    public static String dumpPrivateKey(
            PrivateKey privateKey

    ) {
        notNull(privateKey, "privateKey");

        final byte[] encodedKey = privateKey.getEncoded();
        return "---- Private Key ----\n" +
                "algorithm: " + privateKey.getAlgorithm() + "\n" +
                "format: " + privateKey.getFormat() + "\n" +
                "size: " + encodedKey.length + "\n" +
                "encoded: " + Base64.getEncoder().encodeToString(encodedKey);
    }

    public static void main(
            String[] args
    )
            throws Exception {
        final Scanner scanner = new Scanner(System.in);

        String fileName = "";
        if (args.length > 0) {
            fileName = trimToEmpty(args[0]);
        }

        while (fileName.isEmpty()) {
            System.out.print("input full certificate file path: ");
            fileName = trimToEmpty(scanner.nextLine());
        }

        final String ext = trimToEmpty(getExtension(fileName).toLowerCase());
        final StringBuilder result = new StringBuilder();
        switch (ext) {
            case "crt":
                // 公钥
                try (final InputStream stream = openInputStream(new File(fileName))) {
                    result.append(dumpPublicKey(exportX509Keys(stream)));
                }
                break;
            case "p12":
                // 私钥
                String password = "";
                if (args.length > 1) {
                    password = trimToEmpty(args[1]);
                }

                while (password.isEmpty()) {
                    System.out.print("input password: ");
                    password = trimToEmpty(scanner.nextLine());
                }

                try (final InputStream stream = openInputStream(new File(fileName))) {
                    final Certificate certificate = exportCertificate(stream, password.toCharArray());
                    result.append(dumpCertificate(certificate)).append('\n');
                }

                try (final InputStream stream = openInputStream(new File(fileName))) {
                    final KeyPair keyPair = exportPCKS12Keys(stream, password.toCharArray());
                    result.append(dumpPublicKey(keyPair.getPublic()))
                          .append('\n')
                          .append(dumpPrivateKey(keyPair.getPrivate()));
                }
                break;
            default:
                System.err.print("unknown cert file extension: " + ext);
                return;
        }

        System.out.println(result);
    }
}
